import os
import pandas as pd
from typing import Union, List
from tqdm import tqdm
from dataclasses import dataclass
import uuid


from utils import get_logger
from settings import KW
from utils import get_last_folder_name, save_obj, load_obj


def is_valid_address(
    file_name,
    prov: Union[List[str], str] = None,
    city: Union[List[str], str] = None,
) -> bool:
    """
    根据地址信息，筛选企业
    """

    df = pd.read_csv(file_name, nrows=1)
    try:
        if prov:
            if isinstance(prov, str):
                prov = [prov]

            sign = any(p in df.iloc[0]["所属省份"] for p in prov)
            if not sign:
                return False

        if city:
            if isinstance(city, str):
                city = [city]

            sign = any(c in df.iloc[0]["所属城市"] for c in city)
            if not sign:
                return False
    except:
        print(file_name)
        print(df.iloc[0])
    return True


tqdm.pandas()


def filter_by_keyword(
    df_: pd.DataFrame, cols: Union[List[str], str], keyword: Union[List[str], str]
) -> pd.DataFrame:
    """
    根据关键字筛选企业
    """
    if isinstance(cols, str):
        cols = [cols]
    if isinstance(keyword, str):
        keyword = [keyword]

    # 使用正则表达式模式匹配所有关键字（可选，但仅当关键字数量很大时可能更有效）
    # pattern = '|'.join(r'\b{}\b'.format(re.escape(kw)) for kw in keyword)

    # 直接在选定的列上应用str.contains，并使用any沿行方向聚合结果
    mask = df_[cols].progress_apply(
        lambda row: row.astype(str).str.contains("|".join(keyword), na=False).any(),
        axis=1,
    )
    return df_[mask]


# 示例用法
# df = pd.DataFrame({...})  # 您的DataFrame
# result = filter_by_keyword(df, 'column_name', 'keyword')
# 或者使用多个列和多个关键字
# result = filter_by_keyword(df, ['col1', 'col2'], ['kw1', 'kw2'])ries)

raw_industry_data = "/media/jie/新加卷/pku_data/industry_csv"


def get_df_by_kw_filter(address: dict, industry_kw: list):
    # 遍历 data_home_folder 下的所有文件，找出csv文件
    series = []
    for file_name in os.listdir(raw_industry_data):
        if file_name.endswith(".csv"):
            name = os.path.join(raw_industry_data, file_name)
            if is_valid_address(name, **address):
                # series.append(pd.read_csv(name, nrows=1))
                series.append(
                    filter_by_keyword(
                        pd.read_csv(name),
                        ["企业名称", "经营范围"],
                        industry_kw,
                    )
                )
                print("文件读取结束：", file_name)

    df = pd.concat(series)
    return df


def export_kw_single_address_csv(
    prov: bool, city: bool, prov_city, kw, output_dir="output", use_uuid=False
):
    """
    prov_city : 一个省或者一个市
    适合为了实现市级、省级单位之间的对比

    kw = {"氢能企业":["氢气", "氢能"]}
    export_kw_single_address_csv(
        prov: bool, city: bool, prov_city, kw, output_dir="output", use_uuid=False
    )
    """
    assert prov & city == False
    assert prov | city == True

    if prov:
        address = {"prov": prov_city}
    else:
        address = {"city": prov_city}

    logger = get_logger(prov_city, "industry_kw.log")

    for industry, v in kw.items():
        df = get_df_by_kw_filter(
            address=address,
            industry_kw=v,
        )
        df = df[df["经营状态"].isin(["存续", "开业", "在业", "正常"])]
        logger.info(f"关键词过滤后，共有 {len(df)} 家{industry}企业")

        if use_uuid:
            length = df.shape[0]
            uuids = [str(uuid.uuid4()) for _ in range(length)]
            df["uuid"] = uuids

        fd = f"{output_dir}/{prov_city}"
        if not os.path.exists(fd):
            os.makedirs(fd)

        df.to_csv(f"{fd}/{industry}.csv", index=False)


def export_kw_address_csv(
    address: dict,
    industry_kw: list,
    output_file: str = None,
    use_uuid=True,
    logger=None,
) -> pd.DataFrame:
    """
    df = get_df_by_kw_filter(
        {"city": "武汉市"},
        ["半导体"],
        output_file="jie/test.csv",
        use_uuid=True
    )
    print(df.head())
    """
    if output_file is not None:
        folder_path = os.path.dirname(output_file)
        # 检查文件夹是否存在，若不存在则创建
        if not os.path.exists(folder_path):
            os.makedirs(folder_path)

    # 遍历 data_home_folder 下的所有文件，找出csv文件
    series = []
    for file_name in os.listdir(raw_industry_data):
        if file_name.endswith(".csv"):
            name = os.path.join(raw_industry_data, file_name)
            if is_valid_address(name, **address):
                # series.append(pd.read_csv(name, nrows=1))
                series.append(
                    filter_by_keyword(
                        pd.read_csv(name),
                        ["企业名称", "经营范围"],
                        industry_kw,
                    )
                )
                print("文件读取结束：", file_name)

    df = pd.concat(series)

    df = df[df["经营状态"].isin(["存续", "开业", "在业", "正常"])]
    if logger:
        logger.info(f"关键词过滤后，共有 {len(df)} 家企业")
    else:
        print(f"关键词过滤后，共有 {len(df)} 家企业")

    if use_uuid:
        length = df.shape[0]
        uuids = [str(uuid.uuid4()) for _ in range(length)]
        df["uuid"] = uuids

    if output_file is not None:
        df.to_csv(output_file, index=False)

    return df


# from utils import get_logger


# logger = get_logger(self.NAME, self.logger_file)


def remove_duplicates_and_summarize(fd, logger=None):
    """
    获取fd文件夹下面所有的csv文件，将其拼接到一起，
    再根据uuid属性列去除重复数据

    logger = get_logger("", "去重_kw_filter.log")

    for area in os.listdir("/home/jie/gitee/pku_industry/kw_filter_example/output/kw"):
        remove_duplicates_and_summarize(
            f"/home/jie/gitee/pku_industry/kw_filter_example/output/kw/{area}", logger
        )
    """
    # log
    log_info = []
    name = get_last_folder_name(fd)

    output_file = os.path.join(fd, "去重汇总.csv")

    if os.path.exists(output_file):
        if logger:
            logger.info("汇总文件已存在")
        else:
            print("汇总文件已存在")
        return

    all_data = []

    # 遍历文件夹，读取所有csv文件
    for filename in os.listdir(fd):
        if filename.endswith(".csv"):
            file_path = os.path.join(fd, filename)
            data = pd.read_csv(file_path)
            all_data.append(data)

    # 将所有数据拼接在一起
    combined_data = pd.concat(all_data, ignore_index=True)
    log_info.append(f"{name}汇总数量，去重前:{combined_data.shape[0]}")

    # 根据uuid去除重复数据
    unique_data = combined_data.drop_duplicates(subset="统一社会信用代码")
    log_info.append(f"{name}汇总数量，去重后:{unique_data.shape[0]}")

    # save
    unique_data.to_csv(output_file)

    for info in log_info:
        if logger:
            logger.info(info)
        else:
            print(info)

    return unique_data


def calculate_keyword_frequency(fd, logger=None):
    """
    base_path = "/home/jie/gitee/pku_industry/kw_filter_example/output/kw"

    for area in os.listdir(base_path):
        area_path = os.path.join(base_path, area)
        calculate_keyword_frequency(area_path)
    """
    from settings import get_kw_init_freq

    kw_freq = get_kw_init_freq()
    for file in os.listdir(fd):
        if not file.endswith(".csv") or "汇总" in file:
            continue
        file_path = os.path.join(fd, file)
        df = pd.read_csv(file_path)  # TODO
        industry_name = file.strip(".csv")
        keywords = list(kw_freq[industry_name].keys())

        # 遍历 DataFrame
        for _, row in df.iterrows():
            content = str(row["企业名称"]) + str(row["经营范围"])
            for keyword in keywords:
                if keyword in content:
                    kw_freq[industry_name][keyword] += 1

    save_obj(kw_freq, os.path.join(fd, "kw_freq.pkl"))


def generate_kw_freq_df(fd):
    """
    fd: /home/jie/gitee/pku_industry/kw_filter_example/output/kw
    依据嵌套关键词，生成每个关键词的命中次数表
    KW = {
        "半导体": [
            "半导体设计",
            "半导体制造",
            "半导体设备",
            ...
            "晶圆代工",
            "芯片制造",
        ],
        "集成电路": [
            "集成电路设计",
            ...
            "IC 封测",
        ],
    }
    """

    for industry_name, keywords in KW.items():
        df = pd.DataFrame({keyword: [] for keyword in ["城市"] + keywords})
        for area in os.listdir(fd):
            if not os.path.isdir(os.path.join(fd, area)):
                continue
            obj = load_obj(os.path.join(fd, area, "kw_freq.pkl"))
            d = obj[industry_name]
            d.update({"城市": area})
            for k, v in d.items():
                d[k] = [v]
            df = pd.concat([df, pd.DataFrame(d)], ignore_index=True)
        df.to_csv(os.path.join(fd, f"{industry_name}关键词计数表.csv"))


if __name__ == "__main__":

    # logger = get_logger("", "关键词命中统计.log")

    # generate_kw_freq_df("/home/jie/gitee/pku_industry/kw_filter_example/output/kw")
    export_kw_single_address_csv(
        prov=True,
        city=False,
        prov_city="浙江",
        kw={"氢能企业": ["氢气", "氢能"]},
        output_dir="output",
        use_uuid=False,
    )
